from PIL import Image
import os.path as osp
import os
import xml.etree.ElementTree as ET
import numpy as np
from tqdm import tqdm

def getSizeFormIQ(iq_path):
    tree = ET.parse(iq_path)
    root = tree.getroot()
    ImageParam = root.find('ImageParam')
    height = ImageParam.find('Height').text
    width = ImageParam.find('Width').text

    PyramidParam = root.find('PyramidParam')
    PContrast = PyramidParam.find('PContrast').text
    DContrast = PyramidParam.find('DContrast').text
    BContrast = PyramidParam.find('BContrast').text
    BrightNess = PyramidParam.find('BrightNess').text

    WLGParam = root.find('WLGParam')
    Gamma = WLGParam.find('Gamma').text

    PlayParam = root.find('PlayParam')
    frameCount = PlayParam.find('Count').text
    return int(height), int(width), int(frameCount), float(PContrast), float(DContrast), float(BContrast), float(BrightNess), float(Gamma)

def read_nth_image(file_path, n, _width, _height, dtype=np.uint16):
    bytes_per_pixel = np.dtype(dtype).itemsize
    image_size_bytes = _width * _height * bytes_per_pixel
    # 计算第 n 个图像的起始位置
    offset = n * image_size_bytes
    # 读取目标图像
    image = np.fromfile(file_path, dtype=dtype, count=_width * _height, offset=offset)
    image = image.reshape((_height, _width))  # 调整形状为 (height, width)
    return image

if __name__ == '__main__':
    base_path = r"C:\Users\15519\Desktop\不同部位测试图像\胸_B54_NOSNR"
    file_list = [osp.join(base_path, name) for name in os.listdir(base_path) if name.endswith(".raw")]

    # 获取并保存缩略图
    for raw_path in file_list:
        iq_path = raw_path.replace("raw", "iq")
        raw_msg = getSizeFormIQ(iq_path)
        height, width, frame_count = raw_msg[0], raw_msg[1], raw_msg[2]

        img = read_nth_image(raw_path, frame_count//2, width, height, dtype=np.uint16)
        Image.fromarray(img).save(raw_path.replace(".raw", f"_{height}.png"))

    # 获取保存所有图片
    # tqdm_bar = tqdm(file_list, total=len(file_list))
    # for raw_path in tqdm_bar:
    #     tqdm_bar.set_description_str(raw_path)
    #     iq_path = raw_path.replace("raw", "iq")
    #     raw_msg = getSizeFormIQ(iq_path)
    #     height, width, frame_count = raw_msg[0], raw_msg[1], raw_msg[2]
    #
    #     imgs = np.fromfile(raw_path, dtype=np.uint16).reshape((-1, height, width))
    #     for frame_id in range(frame_count):
    #         img = imgs[frame_id]
    #         Image.fromarray(img).save(raw_path.replace(".raw", f"_frame{frame_id}.png"))
    #         tqdm_bar.set_postfix_str(f"{frame_id+1}/{frame_count}")







